import maya.cmds as cmds

def fix_mash_arnold_color_ultimate():
    # 1. Get the current selection (Select your MASH_Repro mesh)
    selection = cmds.ls(sl=True)
    if not selection:
        cmds.warning("Please select your MASH Repro mesh first!")
        return
    
    for obj in selection:
        # 2. Find the shape node
        shapes = cmds.listRelatives(obj, shapes=True) or [obj]
        if not shapes:
            continue
        
        shape_node = shapes[0]
        
        # FIX 1: Turn on the Arnold vertex color export attribute
        try:
            cmds.setAttr(f"{shape_node}.aiExportColors", 1)
            print(f"Flipped switch: {shape_node}.aiExportColors = 1")
        except:
            try:
                cmds.setAttr(f"{shape_node}.aiExportVertexColors", True)
                print(f"Flipped legacy switch: {shape_node}.aiExportVertexColors = True")
            except:
                print(f"Warning: Could not find an Arnold export color attribute on {shape_node}.")

        # 3. Check if there's already a shader assigned
        shading_groups = cmds.listConnections(shape_node, type='shadingEngine')
        material = None
        mat_type = None
        
        if shading_groups:
            materials = cmds.listConnections(shading_groups[0] + '.surfaceShader')
            if materials:
                material = materials[0]
                mat_type = cmds.objectType(material)
        
        # FIX 2: If there is no shader, or it's NOT an Arnold shader (like lambert1), make one!
        if not material or mat_type != 'aiStandardSurface':
            print(f"No Arnold shader found on {obj}. Generating a new one...")
            
            # Create standard surface shader
            material = cmds.shadingNode('aiStandardSurface', asShader=True, name=f"{obj}_arnoldMat")
            # Create shading engine group
            sg = cmds.sets(renderable=True, noSurfaceShader=True, empty=True, name=f"{material}SG")
            # Connect shader to group
            cmds.connectAttr(f"{material}.outColor", f"{sg}.surfaceShader", force=True)
            # Assign group to the selected mesh
            cmds.sets(obj, edit=True, forceElement=sg)
            
            print(f"Automatically created and assigned: {material}")
        
        # 4. Create the Arnold User Data Color node
        user_data_node = cmds.shadingNode('aiUserDataColor', asTexture=True, name=f"{obj}_MASH_Color")
        
        # Set the attribute name to look for MASH Repro's colorSet
        cmds.setAttr(f"{user_data_node}.colorAttrName", "colorSet", type="string")
        
        # 5. Connect the color node directly to the Base Color of the Arnold Shader
        cmds.connectAttr(f"{user_data_node}.outColor", f"{material}.baseColor", force=True)
        print(f"Success! Connected MASH color to {material}.baseColor via {user_data_node}.\n")

# Run the ultimate tool
fix_mash_arnold_color_ultimate()